package com.analysis.behaviorlogtracing.plugin

import com.android.build.api.transform.Context
import com.android.build.api.transform.DirectoryInput
import com.android.build.api.transform.Format
import com.android.build.api.transform.JarInput
import com.android.build.api.transform.QualifiedContent
import com.android.build.api.transform.Transform
import com.android.build.api.transform.TransformException
import com.android.build.api.transform.TransformInput
import com.android.build.api.transform.TransformInvocation
import com.android.build.api.transform.TransformOutputProvider
import com.android.build.gradle.internal.pipeline.TransformManager
import groovy.io.FileType
import org.apache.commons.codec.digest.DigestUtils
import org.apache.commons.io.FileUtils
import org.apache.commons.io.IOUtils
import java.util.jar.JarEntry
import java.util.jar.JarFile
import java.util.jar.JarOutputStream

class BehaviorLogTracingTransform extends Transform {

    @Override
    String getName() {
        return "Behavior Log Tracing"
    }

    @Override
    Set<QualifiedContent.ContentType> getInputTypes() {
        return TransformManager.CONTENT_CLASS
    }

    @Override
    Set<? super QualifiedContent.Scope> getScopes() {
        return TransformManager.SCOPE_FULL_PROJECT
    }

    @Override
    boolean isIncremental() {
        return false
    }

    @Override
    void transform(TransformInvocation transformInvocation) throws TransformException, InterruptedException, IOException {
        super.transform(transformInvocation)

        transformInvocation.outputProvider.deleteAll()

        transformInvocation.inputs.each { TransformInput input ->
            input.jarInputs.each { JarInput jarInput ->
                traverseJar(jarInput, transformInvocation.outputProvider, transformInvocation.context)
            }

            input.directoryInputs.each { DirectoryInput directoryInput ->
                traverseDir(directoryInput, transformInvocation.outputProvider, transformInvocation.context)
            }
        }

    }

    private static void traverseDir(DirectoryInput directoryInput, TransformOutputProvider outputProvider, Context context) {
        File dir = directoryInput.file
        File dest = outputProvider.getContentLocation(directoryInput.getName(),
                directoryInput.getContentTypes(), directoryInput.getScopes(),
                Format.DIRECTORY)
        FileUtils.forceMkdir(dest)

        String srcDirPath = dir.absolutePath
        String destDirPath = dest.absolutePath

        FileUtils.copyDirectory(dir, dest)

        println("class文件夹名称：" + dir.name + ", size:" + dir.size())

        dir.traverse(type: FileType.FILES, nameFilter: ~/.*\.class/) {
            File inputFile ->
                File edited = editClassFile(dir, inputFile, context.getTemporaryDir())
                if (edited != null) {
                    File target = new File(inputFile.absolutePath.replace(srcDirPath, destDirPath))
                    if (target.exists()) {
                        target.delete()
                    }
                    FileUtils.copyFile(edited, target)
                    edited.delete()
                }
        }
    }

    private static File editClassFile(File dir, File classFile, File tempDir) {
        File EditedFile = null
        FileOutputStream outputStream = null
        try {
            String className = classFile.absolutePath
                    .replace(dir.absolutePath + File.separator, "")
                    .replace(File.separator, ".")
                    .replace(".class", "")
            if (ClassUtil.isDirectoryClassEdit(className)) {
                byte[] sourceClassBytes = BehaviorLogUtil.inputToByteSource(new FileInputStream(classFile))
                byte[] modifiedClassBytes = ASMUtil.editClass(sourceClassBytes)
                if (modifiedClassBytes) {
                    EditedFile = new File(tempDir, UUID.randomUUID().toString() + '.class')
                    if (EditedFile.exists()) {
                        EditedFile.delete()
                    }
                    EditedFile.createNewFile()
                    outputStream = new FileOutputStream(EditedFile)
                    outputStream.write(modifiedClassBytes)
                }
            } else {
                return classFile
            }
        } catch (Exception e) {
            e.printStackTrace()
        } finally {
            IOUtils.closeQuietly(outputStream)
        }
        return EditedFile
    }

    private void traverseJar(JarInput jarInput, TransformOutputProvider outputProvider, Context context) {
        String destName = jarInput.file.name
        def hexName = DigestUtils.md5Hex(jarInput.file.absolutePath).substring(0, 8)
        if (destName.endsWith(".jar")) {
            destName = destName.substring(0, destName.length() - 4)
        }
        File destFile = outputProvider.getContentLocation(
                destName + "_" + hexName, jarInput.contentTypes, jarInput.scopes, Format.JAR
        )
        File editedJar = editJar(jarInput.file, context.getTemporaryDir())
        if (jarInput.file == null || editedJar == null) {
            editedJar = jarInput.file
        }
        FileUtils.copyFile(editedJar, destFile)
    }

    private File editJar(File jarFile, File tempDir) {
        if (jarFile == null || jarFile.size() == 0) {
            return null
        }
        JarFile srcFile = new JarFile(jarFile, false)
        String hexName = DigestUtils.md5Hex(jarFile.absolutePath).substring(0, 8)
        File outputJar =  new File(tempDir, hexName + jarFile.name)
        JarOutputStream jarOutputStream = new JarOutputStream(new FileOutputStream(outputJar))
        Enumeration enumeration = srcFile.entries()
        while (enumeration.hasMoreElements()) {
            JarEntry jarEntry = (JarEntry) enumeration.nextElement()
            InputStream inputStream
            try {
                inputStream = srcFile.getInputStream(jarEntry)
            }catch(Exception e) {
                if (inputStream != null) {
                    IOUtils.closeQuietly(inputStream)
                }
                println("获取jarEntry inputStream 报错 " + jarEntry.getName())
                e.printStackTrace()
                return null
            }
            String entryName = jarEntry.getName()
            if (entryName.endsWith(".DSA") || entryName.endsWith(".SF")) {
                println("不处理 DSA, SF 文件：" + entryName)
                //ignore
            }else {
                String className
                JarEntry outEntry = new JarEntry(entryName)
                byte[] editClassBytes = null
                byte[] sourceClassBytes = null
                try {
                    jarOutputStream.putNextEntry(outEntry)
                    sourceClassBytes = BehaviorLogUtil.inputToByteSource(inputStream)
                }catch(Exception e) {
                    IOUtils.closeQuietly(file)
                    IOUtils.closeQuietly(jarOutputStream)
                    e.printStackTrace()
                    return null
                }
                if (!jarEntry.isDirectory() && entryName.endsWith(".class")) {
                    className = entryName.replace("/", ".").replace(".class", "")
                    if (ClassUtil.isJarClassEdit(className)) {
                        editClassBytes = ASMUtil.editClass(sourceClassBytes)
                    }
                }
                if (editClassBytes == null) {
                    jarOutputStream.write(sourceClassBytes)
                }else {
                    jarOutputStream.write(editClassBytes)
                }
                jarOutputStream.closeEntry()
            }
        }
        jarOutputStream.close()
        srcFile.close()
        return outputJar
    }
}